import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""
import time
import cv2
from NNUnetV1Helper import NNUnetV1Helper

os.makedirs("test", exist_ok=True)

v1_helper = NNUnetV1Helper("../resource/Task670_conus_resized_2000/nnUNetTrainerV2__nnUNetPlansv2.1",
                           fold_num=0)

for each_file in os.listdir("../resource/raw_resize"):
    file_path = os.path.join("../resource/raw_resize", each_file)
    start_time = time.perf_counter()
    out = v1_helper.infer(file_path)
    cv2.imwrite(os.path.join("test", each_file), out)
    end_time = time.perf_counter()
    print("time cost", end_time - start_time, "s")